﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using NUnit.Framework;
using Rhino.Mocks;
using NSsh.Server.TransportLayer;
using NSsh.Server.TransportLayer.State;
using NSsh.Common.Tests;
using NSsh.Common;
using NSsh.Common.Packets;
using NSsh.Common.Types;
using NSsh.Common.Utility;
using Mono.Math;

namespace NSsh.Server.Tests.TransportLayer.State
{
    [TestFixture]
    public class VersionsExchangedStateFixture
    {
        private MockRepository _mockRepository;
        private ITransportLayerManager _manager;
        private KexInitPacket _clientKexInit;
        private KexDHInitPacket _kexDHInit;

        [SetUp]
        public void Setup()
        {
            _mockRepository = new MockRepository();
            _manager = _mockRepository.CreateMock<ITransportLayerManager>();

            Dependency.Clear();
            Dependency.RegisterInstance<ISecureRandom>("SecureRandom", new SecureRandom());

            _clientKexInit = new KexInitPacket();
            _clientKexInit.KexAlgorithms.Names.Add("diffie-hellman-group1-sha1");
            _clientKexInit.ServerHostKeyAlgorithms.Names.Add(PublicKeyAlgorithmHelper.ToString(PublicKeyAlgorithm.DSA));
            _clientKexInit.EncryptionAlgorithmsClientToServer.Names.Add(EncryptionAlgorithmHelper.ToString(EncryptionAlgorithm.Aes128Cbc));
            _clientKexInit.EncryptionAlgorithmsServerToClient.Names.Add(EncryptionAlgorithmHelper.ToString(EncryptionAlgorithm.Aes128Cbc));
            _clientKexInit.MacAlgorithmsClientToServer.Names.Add(MacAlgorithmHelper.ToString(MacAlgorithm.HmacSha1));
            _clientKexInit.MacAlgorithmsServerToClient.Names.Add(MacAlgorithmHelper.ToString(MacAlgorithm.HmacSha1));
            _clientKexInit.CompressionAlgorithmsClientToServer.Names.Add(CompressionAlgorithmHelper.ToString(CompressionAlgorithm.None));
            _clientKexInit.CompressionAlgorithmsServerToClient.Names.Add(CompressionAlgorithmHelper.ToString(CompressionAlgorithm.None));

            _kexDHInit = new KexDHInitPacket();
            _kexDHInit.E = new BigInteger(12345);
        }

        [Test]
        public void ShouldProcessKexCorrectly()
        {
            using (_mockRepository.Record())
            {
                _manager.WritePacket(null);
                LastCall.IgnoreArguments();

                _manager.ServerKexInit = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(null);

                Expect.Call(_manager.ReadPacket()).Return(_clientKexInit);

                _manager.ClientKexInit = _clientKexInit;

                _manager.Parameters = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(new TransportLayerParameters()).Repeat.Twice();
                Expect.Call(_manager.E).Return(null);

                Expect.Call(_manager.ReadPacket()).Return(_kexDHInit);
                _manager.E = _kexDHInit.E;
                Expect.Call(_manager.E).Return(_kexDHInit.E);


                Expect.Call(_manager.E).Return(_kexDHInit.E);
            }

            using (_mockRepository.Playback())
            {
                AbstractTransportState abstractState = new VersionsExchangedState();
                abstractState.Process(_manager);
            }
        }

        [Test]
        public void ShouldFailForBadKexInitAlgorithm()
        {
            _clientKexInit.KexAlgorithms.Names.Clear();
            _clientKexInit.KexAlgorithms.Names.Add("bogus");

            using (_mockRepository.Record())
            {
                _manager.WritePacket(null);
                LastCall.IgnoreArguments();

                _manager.ServerKexInit = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(null);

                Expect.Call(_manager.ReadPacket()).Return(_clientKexInit);

                _manager.ClientKexInit = _clientKexInit;
            }

            using (_mockRepository.Playback())
            {
                Assert.That(new MethodInvoker(delegate
                {
                    AbstractTransportState abstractState = new VersionsExchangedState();
                    abstractState.Process(_manager);

                }), Throws.Exception<InvalidOperationException>());
            }
        }

        [Test]
        public void ShouldFailForBadServerHostKeyAlgorithm()
        {
            _clientKexInit.ServerHostKeyAlgorithms.Names.Clear();
            _clientKexInit.ServerHostKeyAlgorithms.Names.Add("bogus");

            using (_mockRepository.Record())
            {
                _manager.WritePacket(null);
                LastCall.IgnoreArguments();

                _manager.ServerKexInit = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(null);

                Expect.Call(_manager.ReadPacket()).Return(_clientKexInit);

                _manager.ClientKexInit = _clientKexInit;
            }

            using (_mockRepository.Playback())
            {
                Assert.That(new MethodInvoker(delegate
                {
                    AbstractTransportState abstractState = new VersionsExchangedState();
                    abstractState.Process(_manager);

                }), Throws.Exception<InvalidOperationException>());
            }
        }

        [Test]
        public void ShouldFailForBadEncryption()
        {
            _clientKexInit.EncryptionAlgorithmsClientToServer.Names.Clear();
            _clientKexInit.EncryptionAlgorithmsClientToServer.Names.Add("bogus");

            using (_mockRepository.Record())
            {
                _manager.WritePacket(null);
                LastCall.IgnoreArguments();

                _manager.ServerKexInit = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(null);

                Expect.Call(_manager.ReadPacket()).Return(_clientKexInit);

                _manager.ClientKexInit = _clientKexInit;
            }

            using (_mockRepository.Playback())
            {
                Assert.That(new MethodInvoker(delegate
                {
                    AbstractTransportState abstractState = new VersionsExchangedState();
                    abstractState.Process(_manager);

                }), Throws.Exception<InvalidOperationException>());
            }
        }

        [Test]
        public void ShouldFailForBadMac()
        {
            _clientKexInit.MacAlgorithmsServerToClient.Names.Clear();
            _clientKexInit.MacAlgorithmsServerToClient.Names.Add("bogus");

            using (_mockRepository.Record())
            {
                _manager.WritePacket(null);
                LastCall.IgnoreArguments();

                _manager.ServerKexInit = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(null);

                Expect.Call(_manager.ReadPacket()).Return(_clientKexInit);

                _manager.ClientKexInit = _clientKexInit;
            }

            using (_mockRepository.Playback())
            {
                Assert.That(new MethodInvoker(delegate
                {
                    AbstractTransportState abstractState = new VersionsExchangedState();
                    abstractState.Process(_manager);

                }), Throws.Exception<InvalidOperationException>());
            }
        }

        [Test]
        public void ShouldFailForBadCompression()
        {
            _clientKexInit.CompressionAlgorithmsClientToServer.Names.Clear();
            _clientKexInit.CompressionAlgorithmsClientToServer.Names.Add("bogus");

            using (_mockRepository.Record())
            {
                _manager.WritePacket(null);
                LastCall.IgnoreArguments();

                _manager.ServerKexInit = null;
                LastCall.IgnoreArguments();

                Expect.Call(_manager.Parameters).Return(null);
                
                Expect.Call(_manager.ReadPacket()).Return(_clientKexInit);

                _manager.ClientKexInit = _clientKexInit;
            }

            using (_mockRepository.Playback())
            {
                Assert.That(new MethodInvoker(delegate
                    {
                        AbstractTransportState abstractState = new VersionsExchangedState();
                        abstractState.Process(_manager);

                    }), Throws.Exception<InvalidOperationException>());
            }
        }
    }
}
